import numpy as np


class RQI(object):
    def __init__(self, N=51):
        '''
        :param N(int): 养护路段总长度(Km)，向上取整
        '''

        super().__init__()
        # 路面平整度预测模型参数表
        self.a = np.array([[0.1525, 0.1562, 0.1732, 0.1776],
                      [0.1475, 0.1504, 0.1504, 0.1604],
                      [0.1482, 0.1464, 0.1489, 0.1597],
                      [0.1482, 0.1483, 0.1515, 0.1591],
                      [0.1477, 0.1477, 0.1477, 0.1576],
                      [0.1453, 0.1477, 0.1486, 0.1552],
                      [0.1452, 0.1469, 0.1493, 0.1543],
                      [0.1452, 0.1464, 0.1499, 0.1548]])
        self.b = np.array([[0.1825, 0.3221, 0.4643, 0.7032],
                      [0.1575, 0.2996, 0.4796, 0.7196],
                      [0.1418, 0.2961, 0.4686, 0.5843],
                      [0.1418, -0.2918, 0.4618, 0.7193],
                      [0.1314, 0.2814, 0.4614, 0.7057],
                      [0.1311, 0.2614, 0.4514, 0.7014],
                      [0.1524, 0.2764, 0.4507, 0.7057],
                      [0.1314, 0.2836, 0.4368, 0.6836]])
        self.a0 = np.array([0.026])
        self.a1 = np.array([0.65])

        self.N = N

    def IRI(self, t, j, b):
        '''
        :param t: 养护年
        :param j: 养护类别
        :param b: 路段交通量
        :return:
        '''

        x1 = self.a * t + self.b
        res =np.stack([x1[j, b] for _ in range(self.N)]).squeeze()
        return res

    def iri_new(self, t, js, b):
        '''
        :param t: 养护年
        :param j: 养护类别
        :param b: 路段交通量
        :return:
        '''

        x1 = self.a * t + self.b

        # for jx in js:
        #     xxx = x1[jx, b]

        res =np.stack([x1[:, b][jx] for jx in js]).squeeze()
        return res



    def RQI(self, iri):
        return 100.0 / (1.0 + self.a0 * np.exp(self.a1 * iri))


if __name__ == '__main__':
    T = 5
    iri = RQI()
    js = np.random.randint(0, 8, (10, 51))
    b = np.random.randint(0, 4, (10, 1))

    for t in range(T):
        xx = iri.iri_new(t, js, b)
        print(xx)
        rqi = iri.RQI(xx)
        print(rqi)
